from utils.csv_hanlder import select_and_read_files
import matplotlib.pyplot as plt
import numpy as np

def calculate_twiss_parameters(x, x_prime):
    """
    Calculate Twiss parameters (β, α, γ) from phase space coordinates.
    """
    # Remove means to center the phase space ellipse
    x_centered = x - np.mean(x)
    x_prime_centered = x_prime - np.mean(x_prime)
    
    # Calculate second order moments
    sigma_11 = np.mean(x_centered**2)
    sigma_12 = np.mean(x_centered * x_prime_centered)
    sigma_22 = np.mean(x_prime_centered**2)
    
    # Calculate emittance
    emittance = np.sqrt(sigma_11 * sigma_22 - sigma_12**2)
    
    if emittance > 1e-15:
        beta = sigma_11 / emittance
        alpha = -sigma_12 / emittance
        gamma = sigma_22 / emittance
        
        # Verify consistency
        gamma_expected = (1 + alpha**2) / beta
        consistency = np.abs(gamma - gamma_expected) / gamma_expected
        
        if consistency > 0.01:
            print(f"Warning: Twiss parameter consistency check failed: {consistency:.4f}")
    else:
        beta, alpha, gamma = 1.0, 0.0, 1.0
    
    return beta, alpha, gamma, emittance

def plot_phase_space(dfs, filenames=None, save_path='phase_space_3x2.png', figsize=(18, 15)):
    """
    Plot 3x2 grid of phase space diagrams with optimized output formatting.
    """
    fig, axes = plt.subplots(3, 2, figsize=figsize)
    axes = axes.flatten()
    
    # 优化后的表格头
    stats_header = (
        f"{'File':<20} {'f [Hz]':<10} "
        f"{'x_mean':>12} {'x_std':>12} {'x_min':>12} {'x_max':>12} "
        f"{'xp_mean':>12} {'xp_std':>12} {'xp_min':>12} {'xp_max':>12} "
        f"{'y_mean':>12} {'y_std':>12} {'y_min':>12} {'y_max':>12} "
        f"{'yp_mean':>12} {'yp_std':>12} {'yp_min':>12} {'yp_max':>12} "
        f"{'vz_mean':>12} {'vz_std':>12} {'vz_min':>12} {'vz_max':>12}"
    )
    
    twiss_header = (
        f"{'f [Hz]':<12} {'Plane':<6} {'β':>14} {'α':>14} {'γ':>14} {'ε':>14} {'Status':<12}"
    )
    
    # 初始化收集器
    stats_lines = []
    twiss_lines = []
    ellipse_info = []
    
    for idx, df in enumerate(dfs):
        x = df['x']
        y = df['y']
        vx = df['vx']
        vy = df['vy']
        vz = df['vz']
        f_val = df['f'].iloc[0]

        # 计算 x', y' (单位：rad)
        vz_nonzero = np.where(vz != 0, vz, 1e-10)
        x_prime = vx / vz_nonzero
        y_prime = vy / vz_nonzero
        
        # 转换为mm·rad单位（乘以1000：1 rad = 1000 mm·rad）
        x_prime_mmrad = x_prime * 1000.0
        y_prime_mmrad = y_prime * 1000.0

        # 计算Twiss参数（使用原始的rad单位）
        beta_x, alpha_x, gamma_x, emittance_x = calculate_twiss_parameters(x, x_prime)
        beta_y, alpha_y, gamma_y, emittance_y = calculate_twiss_parameters(y, y_prime)
        
        # 更准确的收敛状态判断
        def get_convergence_status(alpha):
            if alpha > 0.01:  # 添加小的容差阈值
                return "Converging"
            elif alpha < -0.01:
                return "Diverging"
            else:
                return "Waist"
        
        status_x = get_convergence_status(alpha_x)
        status_y = get_convergence_status(alpha_y)
        
        # 计算椭圆参数
        ellipse_semi_major_x = np.sqrt(emittance_x * beta_x)
        ellipse_semi_minor_x = np.sqrt(emittance_x / beta_x)
        tilt_angle_x = 0.5 * np.arctan2(2 * alpha_x, gamma_x - beta_x) if abs(alpha_x) > 1e-10 else 0.0
        
        ellipse_semi_major_y = np.sqrt(emittance_y * beta_y)
        ellipse_semi_minor_y = np.sqrt(emittance_y / beta_y)
        tilt_angle_y = 0.5 * np.arctan2(2 * alpha_y, gamma_y - beta_y) if abs(alpha_y) > 1e-10 else 0.0

        # 收集椭圆信息
        ellipse_info.append({
            'file': f"File {idx+1}",
            'freq': f_val,
            'x': {
                'equation': f"γx² + 2αxx' + βx'² = ε: {gamma_x:.2e}x² + {2*alpha_x:.2e}xx' + {beta_x:.2e}x'² = {emittance_x:.2e}",
                'semi_axes': f"a={ellipse_semi_major_x:.2e}, b={ellipse_semi_minor_x:.2e}",
                'tilt': f"θ={np.degrees(tilt_angle_x):.1f}°"
            },
            'y': {
                'equation': f"γx² + 2αxx' + βx'² = ε: {gamma_y:.2e}x² + {2*alpha_y:.2e}xx' + {beta_y:.2e}x'² = {emittance_y:.2e}",
                'semi_axes': f"a={ellipse_semi_major_y:.2e}, b={ellipse_semi_minor_y:.2e}",
                'tilt': f"θ={np.degrees(tilt_angle_y):.1f}°"
            }
        })

        # 绘图 - x-x' 相空间
        # 计算中心点和范围（相对于0）
        x_mean = x.mean()
        x_std = x.std()
        x_prime_mean = x_prime_mmrad.mean()  # 使用mm·rad单位的平均值
        x_prime_std = x_prime_mmrad.std()
        
        # 中心化数据（相对于0），x'使用mm·rad单位
        x_centered = x - x_mean
        x_prime_centered_mmrad = x_prime_mmrad - x_prime_mean
        
        # 确定坐标轴范围，以0为中心
        x_range = max(abs(x_centered.min()), abs(x_centered.max())) * 1.1
        x_prime_range = max(abs(x_prime_centered_mmrad.min()), abs(x_prime_centered_mmrad.max())) * 1.1
        
        sc1 = axes[idx*2].scatter(x_centered, x_prime_centered_mmrad, c=x_prime_centered_mmrad, cmap='viridis', alpha=0.6, s=20)
        axes[idx*2].set_xlim(-x_range, x_range)
        axes[idx*2].set_ylim(-x_prime_range, x_prime_range)  # y轴范围使用mm·rad单位
        
        # 在0处绘制参考线
        axes[idx*2].axhline(y=0, color='black', linestyle='--', linewidth=0.8, alpha=0.5)
        axes[idx*2].axvline(x=0, color='black', linestyle='--', linewidth=0.8, alpha=0.5)
        
        axes[idx*2].set_xlabel('x [mm]')
        axes[idx*2].set_ylabel("x' [mm·rad]")
        title_x = f'File {idx+1}: x-x\' (f={f_val} Hz)\nβ={beta_x:.2e}, α={alpha_x:.2e}, γ={gamma_x:.2e}, ε={emittance_x:.2e}'
        axes[idx*2].set_title(title_x)
        axes[idx*2].grid(True, alpha=0.3)
        plt.colorbar(sc1, ax=axes[idx*2], label="x' [mm·rad]")
        
        # 显示中心点偏移信息（使用mm·rad单位）
        offset_text_x = f'Center: x={x_mean:+.2f} mm, x\'={x_prime_mean:+.2f} mm·rad'
        axes[idx*2].text(0.05, 0.95, f'{status_x}\n{offset_text_x}', 
                        transform=axes[idx*2].transAxes, fontsize=10,
                        verticalalignment='top', bbox=dict(boxstyle='round', 
                        facecolor='lightblue' if status_x == "Converging" else 
                                 'lightcoral' if status_x == "Diverging" else 'lightyellow', 
                        alpha=0.8))

        # 绘图 - y-y' 相空间
        y_mean = y.mean()
        y_std = y.std()
        y_prime_mean = y_prime_mmrad.mean()  # 使用mm·rad单位的平均值
        y_prime_std = y_prime_mmrad.std()
        
        # 中心化数据（相对于0），y'使用mm·rad单位
        y_centered = y - y_mean
        y_prime_centered_mmrad = y_prime_mmrad - y_prime_mean
        
        # 确定坐标轴范围，以0为中心
        y_range = max(abs(y_centered.min()), abs(y_centered.max())) * 1.1
        y_prime_range = max(abs(y_prime_centered_mmrad.min()), abs(y_prime_centered_mmrad.max())) * 1.1
        
        sc2 = axes[idx*2+1].scatter(y_centered, y_prime_centered_mmrad, c=y_prime_centered_mmrad, cmap='plasma', alpha=0.6, s=20)
        axes[idx*2+1].set_xlim(-y_range, y_range)
        axes[idx*2+1].set_ylim(-y_prime_range, y_prime_range)  # y轴范围使用mm·rad单位
        
        # 在0处绘制参考线
        axes[idx*2+1].axhline(y=0, color='black', linestyle='--', linewidth=0.8, alpha=0.5)
        axes[idx*2+1].axvline(x=0, color='black', linestyle='--', linewidth=0.8, alpha=0.5)
        
        axes[idx*2+1].set_xlabel('y [mm]')
        axes[idx*2+1].set_ylabel("y' [mm·rad]")
        title_y = f'File {idx+1}: y-y\' (f={f_val} Hz)\nβ={beta_y:.2e}, α={alpha_y:.2e}, γ={gamma_y:.2e}, ε={emittance_y:.2e}'
        axes[idx*2+1].set_title(title_y)
        axes[idx*2+1].grid(True, alpha=0.3)
        plt.colorbar(sc2, ax=axes[idx*2+1], label="y' [mm·rad]")
        
        # 显示中心点偏移信息（使用mm·rad单位）
        offset_text_y = f'Center: y={y_mean:+.2f} mm, y\'={y_prime_mean:+.2f} mm·rad'
        axes[idx*2+1].text(0.05, 0.95, f'{status_y}\n{offset_text_y}', 
                          transform=axes[idx*2+1].transAxes, fontsize=10,
                          verticalalignment='top', bbox=dict(boxstyle='round', 
                          facecolor='lightblue' if status_y == "Converging" else 
                                   'lightcoral' if status_y == "Diverging" else 'lightyellow', 
                          alpha=0.8))

        # 准备统计表格行
        file_name = filenames[idx] if filenames else f'File {idx+1}'
        file_display = file_name[:20]  # 限制文件名长度
        
        stats_line = (
            f"{file_display:<20} {f_val:<10.0f} "
            f"{x.mean():>12.2e} {x.std():>12.2e} {x.min():>12.2e} {x.max():>12.2e} "
            f"{x_prime_mmrad.mean():>12.2e} {x_prime_mmrad.std():>12.2e} {x_prime_mmrad.min():>12.2e} {x_prime_mmrad.max():>12.2e} "
            f"{y.mean():>12.2e} {y.std():>12.2e} {y.min():>12.2e} {y.max():>12.2e} "
            f"{y_prime_mmrad.mean():>12.2e} {y_prime_mmrad.std():>12.2e} {y_prime_mmrad.min():>12.2e} {y_prime_mmrad.max():>12.2e} "
            f"{vz.mean():>12.2e} {vz.std():>12.2e} {vz.min():>12.2e} {vz.max():>12.2e}"
        )
        stats_lines.append(stats_line)
        twiss_line_x = f"{f_val:<12.0f} {'x':<6} {beta_x:>14.2e} {alpha_x:>14.2e} {gamma_x:>14.2e} {emittance_x:>14.2e} {status_x:<12}"
        twiss_line_y = f"{f_val:<12.0f} {'y':<6} {beta_y:>14.2e} {alpha_y:>14.2e} {gamma_y:>14.2e} {emittance_y:>14.2e} {status_y:<12}"
        twiss_lines.extend([twiss_line_x, twiss_line_y])
    
    print("\n" + "="*180)
    print("PHASE SPACE STATISTICS")
    print("="*180)
    print(stats_header)
    print("-"*180)
    for line in stats_lines:
        print(line)
    print("="*180)
    
    print("\n" + "="*100)
    print("PHASE SPACE ELLIPSE PARAMETERS")
    print("="*100)
    for info in ellipse_info:
        print(f"\n{info['file']} (f={info['freq']} Hz):")
        print(f"  X-plane: {info['x']['equation']}")
        print(f"           Semi-axes: {info['x']['semi_axes']}, Tilt: {info['x']['tilt']}")
        print(f"  Y-plane: {info['y']['equation']}")
        print(f"           Semi-axes: {info['y']['semi_axes']}, Tilt: {info['y']['tilt']}")
    print("="*100)
    
    print("\n" + "="*100)
    print("TWISS PARAMETERS ANALYSIS")
    print("="*100)
    print(twiss_header)
    print("-"*100)
    for line in twiss_lines:
        print(line)
    print("="*100)
    
    # 修复subplot布局问题：增加垂直间距避免标题和xlabel重叠
    plt.tight_layout(rect=[0, 0, 1, 0.96])  # 留出顶部空间给标题
    plt.subplots_adjust(hspace=0.4, wspace=0.3)  # 增加垂直间距(hspace)和水平间距(wspace)
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"\nPhase space plot saved to: {save_path}")
    plt.show()

if __name__ == "__main__":
    # Recommand to select 3 csv files
    dfs, filenames = select_and_read_files(file_type='csv')
    plot_phase_space(dfs, filenames)